## Define workhorse function for running DLMs
run.kalman <- function(y,x) { # Main function
  # First argument is a univariate time series object of length T with no missing data
  # Second argument is a table time series object of dimension T x k with no missing data
  
  # Helper function to build DLM i.e. regression with time-varying parameters
  buildTVP <- function(u) { 
    try(dlmModReg(x, addInt = F, dV = exp(u[1]),
                  dW = exp(u[2:num]), 
                  m0 = rep(0, (num-1)),
                  C0 = 1e+07 * diag((num-1))
    ), silent=F)
  }
  
  num <- NCOL(x)+1 #pick right dimension of starting values for size of x matrix
  outMLE <- try(dlmMLE(y, parm = rep(0, (num)), 
                       buildTVP, hessian = TRUE, 
                       method = "BFGS", debug = TRUE)) #Estimate starting variance & means
  
  if(outMLE$convergence) {
    return(NULL) # 
  } else if (!exists("outMLE")) {
    return(NULL)
  } else {
    mod <- buildTVP(outMLE$par) # Build smoothing model
    filtered <- dlmFilter(y, mod) # Run Kalman filter
    out <- dlmSmooth(filtered) # Run smoother on filtered model
    return(list(out, mod, filtered)) #Return smoother results and the logical vector for selecting cases
  }
}


## Define function to prep dlm results for graphing coefficients
graph.DLM.prep <- function(mod, year){ 
  # NB: This function uses the data.table package
  # First argument is an estimated smoothed model (from run.kalman)
  # Second argument is the year labels from the original data
  # (NB: code as written only works for annual data)
  year <- c(year, (max(year) + 1))
  xaxis <- year
  
  #Axis label as years ending in 0 for plot axis labels
  xaxis[!grepl("[[:digit:]]{3}0", xaxis)] <- "" 
  
  param <- mod$s
  colnames(param) <- c(paste0("beta", 1:ncol(param)))
  
  #Storing parameter labels for matrix to send to plot
  labels <- rep(colnames(param), nrow(param)) 
  
  #Assemble and rename data
  to.graf <- as.data.frame(cbind(labels[order(labels)], 
                                 as.vector(param),
                                 rep(year, ncol(param))),
                           stringsAsFactors = FALSE)
  to.graf[, 2] <- as.numeric(as.character(to.graf[, 2]))
  names(to.graf) <- c("param", "value", "year") #Column names for matrix to plot
  
  ##Get confidence intervals
  vc.matrices <- dlmSvd2var(mod$U.S, mod$D.S) #Calculate VC matrix of parameters at each discrete time point
  intervals <- list() #Initialize empty list for confidence interval dataframes
  for (i in 1:nrow(param)) { #Extract confidence intervals for coefficient estimates at each discrete time point
    intervals[[i]] <- as.data.frame(cbind(colnames(param),
                                          param[i,] - 1.96 * sqrt(diag(vc.matrices[[i]])),
                                          param[i,] + 1.96 * sqrt(diag(vc.matrices[[i]])),
                                          rep(year[i], ncol(param))))
  }
  cis <- as.data.frame(rbindlist(intervals)) 
  
  #Next few lines just munge the confidence intervals into a plot-ready format
  cis[, 2] <- as.numeric(as.character(cis[, 2]))
  cis[, 3] <- as.numeric(as.character(cis[, 3]))
  names(cis) <- c("param", "ci.lo", "ci.hi", "year")
  final.to.graf <- merge(to.graf, cis)
  
  #Return table of coefficient estimates and confidence intervals and custom axis labels
  return(list(final.to.graf, as.numeric(xaxis))) 
}


## Define function for making standardized DLM coefficient graphs
plot.DLM.sep <- function(coef, xaxis, lab, y.hi = NULL, y.lo = NULL) { #Plots a single time-varying coefficient from a DLM
  # First argument is a table of point est. and conf. intervals subset from item one of the list output from "graph.DLM.prep" function
  # Second argument is a numeric vector of axis breaks - i.e. item two of the list output from "graph.DLM.prep" function
  # Third argument is a single string of the name of the variable to serve as the plot title
  # Fourth and fifth arguments are optional vertical plot limits
  
  plot <- ggplot(coef, aes(x = year) ) + 
    geom_line(aes(y = value)) + 
    geom_ribbon(aes(ymin = ci.lo, ymax = ci.hi, alpha = .001)) +
    scale_y_continuous("Coefficient", limits = c(y.lo, y.hi)) + 
    scale_x_continuous("", breaks = xaxis) + 
    theme(panel.grid.major = element_blank(), 
          panel.grid.minor = element_blank(), 
          panel.background = element_blank(), 
          axis.line = element_line(colour = "black"),
          axis.ticks = element_line(colour = "black"),
          axis.text = element_text(size = 8, colour = "black"), 
          axis.title.y=element_text(size = 10, colour = "black")) +
    ggtitle(lab) + 
    geom_hline(yintercept = 0, linetype = "dashed", color = "red") + 
    theme(legend.position = "none")
  
  return(plot)
}
